import numpy as np
from structure_tensor import eig_special_2d, structure_tensor_2d, structure_tensor_3d
import matplotlib.pyplot as plt
# from RootUtils.DataClass.Seismic import Seismic
# from RootUtils.BasicAlgorithm import Interplation as itp
# from RootUtils.Cube import CubePlot


def cal_structure_tensor(cube, sigma, rho, plt_min, plt_max, model='3d'):

    shape = cube.shape
    x_mid = int(shape[0] / 2)
    plt_min= plt_min
    plt_max = plt_max

    if model == '2d':

        # Load 2D data.
        image = cube
        S = structure_tensor_2d(image, sigma, rho)
        return S[0,:,:]
        plt.subplot(1, 7, 1)
        plt.imshow(image, cmap='seismic_r')
        plt.subplot(1, 7, 2)
        plt.imshow(S[0, :, :])
        # plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 3)
        plt.imshow(S[1, :, :])
        # plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 4)
        plt.imshow(S[2, :, :])
        # plt.clim(plt_min, plt_max)
        # plt.subplot(1, 7, 5)
        # plt.imshow(S[3, :, :])
        # # plt.clim(plt_min, plt_max)

    if model == '3d':
        S = structure_tensor_3d(cube, sigma, rho)
        plt.subplot(1, 7, 1)
        plt.imshow(cube[x_mid, :, :].T, cmap='seismic_r')
        plt.subplot(1, 7, 2)
        plt.imshow(S[0, x_mid, :, :].T)
        plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 3)
        plt.imshow(S[1, x_mid, :, :].T)
        plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 4)
        plt.imshow(S[2, x_mid, :, :].T)
        plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 5)
        plt.imshow(S[3, x_mid, :, :].T)
        plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 6)
        plt.imshow(S[4, x_mid, :, :].T)
        plt.clim(plt_min, plt_max)
        plt.subplot(1, 7, 7)
        plt.imshow(S[5, x_mid, :, :].T)
        plt.clim(plt_min, plt_max)

    plt.show()
    # val, vec = eig_special_2d(S)
    return S


if __name__ == '__main__':

    seismic_case = Seismic.Seismic()
    seis_path = r'E:\Lion\Data\Lion-2200-3400-6370-7500-3200-4600.sgy'
    seismic_case.sgy_import(seis_path, iline=9, xline=21, cdpx=73, cdpy=77)
    cube = seismic_case.seisnc['data'].values
    S = cal_structure_tensor(
        cube,
        sigma=1,
        rho=2,
        plt_min=0,
        plt_max=200000,
        model='3d'
    )

    seismic_case.seisnc['data'].values = S[2,:, :, :]
    seis_data_np_path = seismic_case.segy_file_path + '_tensor2.npy'
    seismic_case.segy_export('tensor2')